# encoding=‘utf-8'
import base64
import hashlib
import re

from Crypto.Cipher import AES


def getSha256(text):
    h = hashlib.sha256()
    h.update(text.encode())
    ret = h.hexdigest()
    return ret


class AES_Crypt(object):
    def __init__(self, key, iv):
        self.key = key.encode("utf-8")
        self.iv = (iv[20:36]).encode()
        self.mode = AES.MODE_CBC
        # 加密函数，如果text不足16位就用空格补足为16位，
        # 如果大于16当时不是16的倍数，那就补足为16的倍数。

    def encrypt(self, text):
        try:
            text = text
            cryptor = AES.new(self.key, self.mode, self.iv)
            # 这里密钥key 长度必须为16（AES-128）,
            # 24（AES-192）,或者32 （AES-256）Bytes 长度
            # 目前AES-128 足够目前使用
            length = 16
            count = len(text.encode('utf-8'))
            if count < length:
                add = (length - count)
                # \0 backspace
                # 以\0补足js做空格切割，切不掉，换成ASCII字符
                # text = text + ('\0' * add)
                text = text + (chr(add) * add)
            elif count > length:
                add = (length - (count % length))
                # text = text + ('\0' * add)
                text = text + (chr(add) * add)
            ciphertext = cryptor.encrypt(text.encode("utf-8"))
            # 因为AES加密时候得到的字符串不一定是ascii字符集的，输出到终端或者保存时候可能存在问题
            # 所以这里统一把加密后的字符串转化为字符串
            s = str(base64.b64encode(ciphertext), 'utf-8')
        except Exception:
            s = ""
        return s

        # 解密后，去掉补足的空格用strip() 去掉

    def decrypt(self, text):
        try:
            cryptor = AES.new(self.key, self.mode, self.iv)
            text = base64.b64decode(text)
            plain_text = cryptor.decrypt(text)
            # return plain_text.rstrip('\0')
            # print(plain_text)
            # 去除混乱字符
            reboj = re.compile('[\\x00-\\x08\\x0b-\\x0c\\x0e-\\x1f\n\r\t]')
            return reboj.sub('', plain_text.decode())
            # return bytes.decode(plain_text).rstrip('\0')
        except Exception:
            return ""


def setAES(key, iv, text):
    ac = AES_Crypt(key, iv)  # 初始化密钥
    ret = ac.encrypt(text)
    return ret


def getAES(key, iv, text):
    ac = AES_Crypt(key, iv)  # 初始化密钥
    ret = ac.decrypt(text)
    return ret


def setDataAes(key, text):
    key = getSha256(key)[10:26]
    iv = getSha256(key)
    ret = setAES(key, iv, text)
    return ret


def getDataAes(key, text):
    key = getSha256(key)[10:26]
    iv = getSha256(key)
    ret = getAES(key, iv, text)
    return ret


if __name__ == '__main__':
    ret = setDataAes("Wchime", "minioadmin")
    print(ret)
    print(setDataAes("Wchime", "admin"))


